from IPython.display import clear_output
import random
import numpy as np
import matplotlib.pyplot as plt
import math
import time
import copy
from grid_world_v2 import GridWorld
from replay_memory import ReplayMemory

from keras.models import Sequential
from keras.layers.core import Dense, Dropout, Activation
from keras.optimizers import RMSprop
from keras.optimizers import Adam
from keras.models import model_from_json
from keras.callbacks import Callback
from random import shuffle
from keras.layers import MaxPooling2D, Flatten
from keras.layers import Convolution2D

obstacle_grid = np.array([[0, 0, 0, 1, 1, 1, 1, 0, 0, 0],
                          [0, 0, 0, 1, 1, 1, 1, 0, 0, 0],
                          [0, 0, 0, 1, 1, 1, 1, 0, 0, 0],
                          [0, 0, 0, 1, 1, 1, 1, 0, 0, 0],
                          [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                          [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                          [0, 0, 0, 1, 1, 1, 1, 0, 0, 0],
                          [0, 0, 0, 1, 1, 1, 1, 0, 0, 0],
                          [0, 0, 0, 1, 1, 1, 1, 0, 0, 0],
                          [0, 0, 0, 1, 1, 1, 1, 0, 0, 0]])

GRID_SIZE = 10
TOTAL_CARS = 2
TERMINAL_REWARD = 5
DEMAND_LIMIT = 1

def testNaivePolicy(TOTAL_CARS, GRID_SIZE, TERMINAL_REWARD, DEMAND_LIMIT, obstacle_grid):
    terminal = False
    grid_world = GridWorld(num_cars=TOTAL_CARS, grid_size=GRID_SIZE, terminal_reward=TERMINAL_REWARD,
                           demand_limit=DEMAND_LIMIT, obstacle_grid=obstacle_grid)

    while (not terminal and grid_world.t < 50):
        grid_world.t += 1
        # grid_world.print_grid()
        all_agents_step = grid_world.stepAll('Naive', 1)
        terminal = grid_world.isTerminal()

    return grid_world.t

testNaivePolicy(TOTAL_CARS, GRID_SIZE, TERMINAL_REWARD, DEMAND_LIMIT, obstacle_grid)

moves = 0
games = 2000
for i in range(games):
    print
    'Game: ', i
    moves = moves + testNaivePolicy(TOTAL_CARS, GRID_SIZE, TERMINAL_REWARD, DEMAND_LIMIT, obstacle_grid)
    clear_output(wait=True)
result = moves / float(games)
print
result

### Test Model

model = Sequential()
model.add(Convolution2D(32, 3, 3, input_shape=(4, GRID_SIZE, GRID_SIZE), activation='relu'))
model.add(Convolution2D(32, 3, 3, activation='relu'))
model.add(Convolution2D(32, 3, 3, activation='relu'))
model.add(Convolution2D(32, 3, 3, activation='relu'))
model.add(Flatten())
model.add(Dense(256, activation='relu'))
model.add(Dense(5, init='lecun_uniform'))
model.add(Activation('linear'))
rms = RMSprop()
model.compile(loss="mse", optimizer=rms)

model_json = model.to_json()
with open("models/model_10x10_Conv.json", "w") as json_file:
    json_file.write(model_json)

# load json and create model
json_file = open('models/model_10x10_Conv.json', 'r')
loaded_model_json = json_file.read()
json_file.close()
loaded_model1 = model_from_json(loaded_model_json)

loaded_model1.load_weights("weights/Gridworld_6x6_4Layer_Conv_300.h5")
print("Loaded model from disk")

def testModel(model, TOTAL_CARS, GRID_SIZE, TERMINAL_REWARD, DEMAND_LIMIT, obstacle_grid):
    terminal = False
    grid_world = GridWorld(num_cars=TOTAL_CARS, grid_size=GRID_SIZE, terminal_reward=TERMINAL_REWARD,
                           demand_limit=DEMAND_LIMIT, obstacle_grid=obstacle_grid)

    while (not terminal and grid_world.t < 50):
        grid_world.t += 1
        # grid_world.print_grid()
        all_agents_step = grid_world.stepAll(model, 0)
        terminal = grid_world.isTerminal()

    return grid_world.t

testModel(loaded_model1, TOTAL_CARS, GRID_SIZE, TERMINAL_REWARD, DEMAND_LIMIT, obstacle_grid)

moves = 0
games = 2000
for i in range(games):
    print
    'Game: ', i
    moves = moves + testModel(loaded_model1, TOTAL_CARS, GRID_SIZE, TERMINAL_REWARD, DEMAND_LIMIT, obstacle_grid)
    clear_output(wait=True)
result = moves / float(games)
print(result)

## 7x7

obstacle_grid = np.array([[0, 0, 1, 1, 1, 0, 0],
                          [0, 0, 0, 1, 0, 0, 0],
                          [0, 0, 1, 1, 1, 0, 0],
                          [0, 0, 0, 0, 0, 0, 0],
                          [0, 0, 1, 1, 1, 0, 0],
                          [0, 0, 0, 1, 0, 0, 0],
                          [0, 0, 1, 1, 1, 0, 0]])

GRID_SIZE = 7
TOTAL_CARS = 2
TERMINAL_REWARD = 5
DEMAND_LIMIT = 1
WAIT_TIME = 10

def testNaivePolicy(TOTAL_CARS, GRID_SIZE, TERMINAL_REWARD, DEMAND_LIMIT, WAIT_TIME, obstacle_grid):
    terminal = False
    grid_world = GridWorld(num_cars=TOTAL_CARS, grid_size=GRID_SIZE, terminal_reward=TERMINAL_REWARD,
                           demand_limit=DEMAND_LIMIT, cust_wait_time=WAIT_TIME, obstacle_grid=obstacle_grid)
    while (not terminal and grid_world.t < 50):
        grid_world.t += 1
        grid_world.print_grid()
        all_agents_step = grid_world.stepAll('Naive', 1)
        terminal = grid_world.isTerminal()

    return grid_world.t

testNaivePolicy(TOTAL_CARS, GRID_SIZE, TERMINAL_REWARD, DEMAND_LIMIT, WAIT_TIME, obstacle_grid)

moves = 0
games = 2000
for i in range(games):
    print
    'Game: ', i
    moves = moves + testNaivePolicy(TOTAL_CARS, GRID_SIZE, TERMINAL_REWARD, DEMAND_LIMIT, WAIT_TIME, obstacle_grid)
    clear_output(wait=True)
result = moves / float(games)

model = Sequential()
model.add(Dense(256, init='lecun_uniform', input_shape=(GRID_SIZE * GRID_SIZE * 4,)))
model.add(Activation('relu'))
model.add(Dropout(0.2))
model.add(Dense(256, init='lecun_uniform'))
model.add(Activation('relu'))
model.add(Dropout(0.2))
model.add(Dense(5, init='lecun_uniform'))
model.add(Activation('linear'))
rms = RMSprop()
model.compile(loss="mse", optimizer=rms)

model_json = model.to_json()
with open("models/model_7x7_4layer.json", "w") as json_file:
    json_file.write(model_json)

# load json and create model
json_file = open('models/model_7x7_4layer.json', 'r')
loaded_model_json = json_file.read()
json_file.close()
loaded_model1 = model_from_json(loaded_model_json)

loaded_model1.load_weights("weights/Gridworld_7x7_4Layer_120.h5")
print("Loaded model from disk")


def testModel(model, TOTAL_CARS, GRID_SIZE, TERMINAL_REWARD, DEMAND_LIMIT, WAIT_TIME, obstacle_grid):
    terminal = False
    grid_world = GridWorld(num_cars=TOTAL_CARS, grid_size=GRID_SIZE, terminal_reward=TERMINAL_REWARD,
                           demand_limit=DEMAND_LIMIT, cust_wait_time=WAIT_TIME, obstacle_grid=obstacle_grid)

    while (not terminal and grid_world.t < 50):
        grid_world.t += 1
        grid_world.print_grid()
        all_agents_step = grid_world.stepAll(model, 0)
        terminal = grid_world.isTerminal()

    return grid_world.t

testModel(loaded_model1, TOTAL_CARS, GRID_SIZE, TERMINAL_REWARD, DEMAND_LIMIT, WAIT_TIME, obstacle_grid)

moves = 0
games = 2000
for i in range(games):
    print
    'Game: ', i
    moves = moves + testModel(loaded_model1, TOTAL_CARS, GRID_SIZE, TERMINAL_REWARD, DEMAND_LIMIT, WAIT_TIME,
                              obstacle_grid)
    clear_output(wait=True)
result = moves / float(games)
